clear; clc; close all;



%% Initializing parameters for example Figures 1 and 3
%% for cstar(y) we will use a consumption policy function similar to our application in Section 4

rng('default')


bta =  .95; % time discount factor
abar = 0.25; % ss or initial a
mem_discount = 1; % memory parameter
 
 
r =    1*(1/bta - 1) - 0.005; % interest rate
w = 1.2;
log_c = 0;


b= 0/(1+r); % borrowing constraint
amin = b;
amax = 45;
na   = 1001;

util_type = 'crra';
gam = 2; % parameter of the utility function u(c) = c - 0.5*gamma*c^2
%gam = 2; % parameter of the utility function u(c) = c - 0.5*gamma*c^2

% state y
% y_t = ybar + eps_yt + eps_yit
% eps_yt ~ N(0,sigma_y^2)
% eps_yit ~ N(0,sigma_yi^2)
ybar = 1;
sigma_y = 0*0.2; % 1*0.1; keep total variance of y_t the same as with common shocks
sigma_yi =  0.2;
rho_y = 0; % delete this eventually
ybar_logn =  - 0.5*(sigma_y^2 + sigma_yi^2);


% coefficients on optimal policy in respect to state (a_t + y_t)
cstar_ya = r/(1+r);

% endogenous state a: r fully characterizes LOM for a
% a_{t+1} = (1+r_t)*(y_t + a_t - c_t)

% steady state action
cbar = ybar + r/(1+r)*abar;

% priors
% mu0: generalized prior for chat_0(y,a)
%  prior cov parameters:
%  sigma_0((y,a),(y',a')) = sigma_c^2*exp(-psi*((y-y')^2+(a-a')^2))

prior_mean_coeff = 1;

mu0_fun_c     = @(ya) cbar + prior_mean_coeff*r/(1+r)*(ya - ybar - abar);
mu0_fun_clogn = @(ya) exp(cbar) + prior_mean_coeff*ya - ybar - abar;

W_cc = 1;
kappa =  0.05;
sigma_c = 0.14;
psi = 0.4;


%  covariance function
cov_fun = 1;  % 1 = squared-exponential, 2 = abs value-exponential
ex_signals = 0.0*0.15^2;  % exogenous signals

%  reasoning action
res_action = 1;  % 1 = consumption, 2 = savings (next period wealth)

factor_mpc = 1;

 

[xs,~,cstar,~,ctilde] = bc_opt_pol_iid_egm(bta,      r, b, ybar_logn + log(w),        (sigma_y^2+sigma_yi^2)^0.5,gam, amin,amax,na, util_type,0,0);

[xs,xi] = unique(xs);
temp = griddedInterpolant(xs,cstar(xi),'linear','linear');
cstar = @(x) temp(x);

stateGrid_size = 2001;
grid_min = 0;
grid_max = 10;

%stateGrid = linspace(xs(1), 2*xs(end), stateGrid_size)';
stateGrid = linspace(grid_min, grid_max, stateGrid_size)';
stateGrid_step = stateGrid(2) - stateGrid(1);


temp = griddedInterpolant(xs,ctilde(xi),'linear','linear');
ctilde = @(x) temp(x);
mu0_fun = ctilde;
ctilde = mu0_fun(stateGrid);

rw_pol = (r/(1+r))*stateGrid' + 1/(1+r)*w;


%% First Figure

sigma_c = sqrt(0.03); 
kappa   = 0.01;
psi     = 0.25; 

sigma_etasq = kappa*sigma_c^2/(sigma_c^2 - kappa);


ya_hist = [4.5; 5.5];
sigma_etasq_hist = [ sigma_etasq ; sigma_etasq];
eta_hist = mu0_fun(ya_hist) + sqrt(sigma_etasq_hist).*[1;-1];

[~,ya_ind] = min(abs(stateGrid - ya_hist(1))); 

[mu_t2, Sigma_t2] = gp_update(stateGrid, ya_hist, eta_hist, sigma_etasq_hist, mu0_fun,sigma_c,psi, cov_fun);

ya_hist_p = [3.25;  6.75];
sigma_etasq_hist = 1*[ sigma_etasq ; sigma_etasq];
eta_hist_p = mu0_fun(ya_hist_p) + sqrt(sigma_etasq_hist).*[1;-1];
[~,ya_ind_p] = min(abs(stateGrid - ya_hist_p(1))); 


[mu_t2_p, Sigma_t2_p] = gp_update(stateGrid, ya_hist_p, eta_hist_p, sigma_etasq_hist, mu0_fun,sigma_c,psi, cov_fun);
%[mu_t2, Sigma_t2_highpsi] = gp_update(stateGrid, ya_hist, eta_hist, sigma_etasq_hist, mu0_fun,sigma_c,psi*5, cov_fun);


%%%%%%% Figure I
figure
plot(stateGrid, Sigma_t2,'Linewidth',2.5)
xlim([2,8])
%ylim( [0.006,0.025])
ylim( [0.8*min(Sigma_t2), sigma_c^2*1.2])
hold on
%plot(stateGrid, sigma_c^2*ones(size(stateGrid)),'k--')
plot(stateGrid, Sigma_t2(ya_ind)*ones(size(stateGrid)), 'k:','LineWidth', 1); 
%scatter(ya_hist1, sigma_c^2,150,'k')
vline( ya_hist(1), ':k', '', 16, 1, 2.5)
vline( ya_hist(2), ':k', '', 16, 1, 2.5)
set(gca, 'XTick', ya_hist);
set(gca, 'XTickLabel', {'$y_1$', '$y_2$'}, 'TickLabelInterpreter', 'latex');
set(gca, 'YTick', []);
%set(gca, 'YTickLabel', {'$\sigma_c^2$'}, 'TickLabelInterpreter', 'latex');
ax = gca;
ax.FontSize = 17;
ylabel('$\widehat \sigma_{2}^2(y)$  (Beginning-of-period uncertainty)','interpreter','latex','FontSize',14)
xlabel('State ($y$)','interpreter','latex')
legend('$\widehat \sigma_{2}^2(y)$','interpreter','latex')
hold off
print -depsc example_cor2.eps

figure
plot(stateGrid, Sigma_t2_p,'Linewidth',2.5)
xlim([1,9])
%ylim( [0.006,0.025])
ylim( [0.8*min(Sigma_t2), sigma_c^2*1.2])
hold on
%plot(stateGrid, sigma_c^2*ones(size(stateGrid)),'k--')
plot(stateGrid, Sigma_t2_p(ya_ind_p)*ones(size(stateGrid)), 'k:','LineWidth', 1); 
%scatter(ya_hist1, sigma_c^2,150,'k')
vline( ya_hist_p(1), ':k', '', 16, 1, 2.5)
vline( ya_hist_p(2), ':k', '', 16, 1, 2.5)
set(gca, 'XTick', ya_hist_p);
set(gca, 'XTickLabel', {'$y_1$', '$y_2$'}, 'TickLabelInterpreter', 'latex');
set(gca, 'YTick', []);
%set(gca, 'YTickLabel', {'$\sigma_c^2$'}, 'TickLabelInterpreter', 'latex');
ax = gca;
ax.FontSize = 17;
ylabel('$\widehat \sigma_{2}^2(y)$  (Beginning-of-period uncertainty)','interpreter','latex','FontSize',14)
xlabel('State ($y$)','interpreter','latex')
legend('$\widehat \sigma_{2}^2(y)$','interpreter','latex')
hold off
print -depsc example_cor2_v2.eps


%%


ya_hist1 = 2.5;
sigma_etasq_hist = kappa*sigma_c^2/(sigma_c^2 - kappa);
eta_hist1 = mu0_fun(ya_hist1) + 1*sqrt(sigma_etasq_hist);

[mu_t1, Sigma_t1] = gp_update(stateGrid, ya_hist1, eta_hist1, sigma_etasq_hist, mu0_fun,sigma_c,psi, cov_fun);
[mu_t1_psi0, Sigma_t1_psi0] = gp_update(stateGrid, ya_hist1, eta_hist1, sigma_etasq_hist, mu0_fun,sigma_c,0, cov_fun);

ya_hist = [2.5; 1.5];
sigma_etasq_hist = [ kappa*sigma_c^2/(sigma_c^2 - kappa); kappa*Sigma_t1(301)/(Sigma_t1(301) - kappa)];
eta_hist = mu0_fun(ya_hist) + 1*sqrt(sigma_etasq_hist).*[1;1];

[mu_t2, Sigma_t2] = gp_update(stateGrid, ya_hist, eta_hist, sigma_etasq_hist, mu0_fun,sigma_c,psi, cov_fun);

ya_hist5 = [2.5; 1.5; 4];
sigma_etasq_hist = [ kappa*sigma_c^2/(sigma_c^2 - kappa); kappa*Sigma_t1(301)/(Sigma_t1(301) - kappa); kappa*Sigma_t2(801)/(Sigma_t1(801) - kappa)];
eta_hist = mu0_fun(ya_hist5) + 1*sqrt(sigma_etasq_hist).*[1;1; 0];

[mu_t5, Sigma_t5] = gp_update(stateGrid, ya_hist5, eta_hist, sigma_etasq_hist, mu0_fun,sigma_c,psi, cov_fun);


%%%%%%%%%%%%%%%%%%   Figure II.(a)
figure
plot(stateGrid, sigma_c^2*ones(size(stateGrid)),'k--','Linewidth',2); hold on
plot(stateGrid, [Sigma_t1],'Linewidth',2.5, 'Color',[0 0.4470 0.7410])
plot(stateGrid, [Sigma_t2],'-.','Linewidth',2.5,'Color',[0.8500 0.3250 0.0980])
xlim([0.3,5])
ylim( [0.005,0.04])
%ylim( [-0.000,0.025])
plot(stateGrid, kappa*ones(size(stateGrid)), 'k:','LineWidth', 1.5); 
%scatter( mean(ya_hist), Sigma_t2(401),150,'r', 'filled')
scatter( stateGrid(451), Sigma_t2(451),150,'r', 'filled')
scatter( stateGrid(361), Sigma_t2(361),150,'r', 'filled')
scatter(ya_hist(2),  Sigma_t1(301),150,'b','filled')
scatter(stateGrid(801),  Sigma_t2(801),150,'filled','MarkerFaceColor',[1 0 1])
scatter(stateGrid(501),  sigma_c^2,150,'filled','k')
%scatter(stateGrid(1001),  Sigma_t2(1001),150,'y','filled')
vline( ya_hist(1),      ':k', '', 16, 1, 2.5)
vline( ya_hist(2),      ':k', '', 16, 1, 2.5)
%vline( mean(ya_hist),   ':k', '', 16,1,2.5)
vline( stateGrid(361),   ':k', '', 16,1,2.5)
vline( stateGrid(451),   ':k', '', 16,1,2.5)
vline( stateGrid(801), ':k', '', 16,1,2.5)
xticks( sort([ya_hist5;  stateGrid([361;451])]))
xticklabels({'$y_2$','$y_4$','$y_3$','$y_1$','$y_5$'})
yticks([kappa, sigma_c^2])
yticklabels( {'$\kappa$','$\sigma_c^2$'})
ax = gca;
ax.FontSize = 17;
set(gca,'TickLabelInterpreter', 'latex')
ylabel('System 1 uncertainty ($\widehat \sigma^2$)','interpreter','latex')
xlabel('State ($y$)','interpreter','latex')
hold off
legend('$\widehat \sigma_0^2(y)$','$\widehat \sigma_1^2(y)$',' $\widehat \sigma_2^2(y) = \widehat \sigma_3^2(y) = \widehat \sigma_4^2(y)$','Location','Northeast','interpreter','latex','FontSize',17)
print -depsc example_shat2_v3.eps


%%


%%%%%%%%%%%%%%%%%%   Figure II.(b)

figure
plot((1:5)', [1 - kappa/(sigma_c^2); 1 - kappa/Sigma_t1(301); 0; 0;  1 - kappa/Sigma_t2(801)]','k','Linewidth',1.5); hold on
xlim([0.5,5.5])
ylim( [-0.1,0.85])
plot(0:5, zeros(1,6),'k--')
scatter( 1, 1 - kappa/sigma_c^2,125,'k','filled')
scatter( 2, 1 - kappa/Sigma_t1(301),125,'b','filled')
scatter( 3, 0,125,'r','filled')
scatter( 4, 0,125,'r','filled')
scatter( 5,  1 - kappa/Sigma_t2(801),125,'filled','MarkerFaceColor',[1 0 1])
xticks( [1:5] )
yticks([0; 1 - kappa/sigma_c^2])
yticklabels( {'0','$\frac{\sigma_c^2 - \kappa}{\sigma_c^2}$'})
ax = gca;
ax.FontSize = 17;
set(gca,'TickLabelInterpreter', 'latex')
ylabel('System 2 revision weight ($\alpha^*$)','interpreter','latex')
xlabel('Time period ($t$)','interpreter','latex')
hold off
legend('$\alpha_t^*(y_t)$','Location','Northeast','interpreter','latex','FontSize',17)
print -depsc example_alphat_star_v3.eps

%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%

%%%%%%%%%%%%% Figure B.1
 
figure
plot(stateGrid, [mu_t1],'k-.','LineWidth',2.5); hold on
plot(stateGrid, [mu_t2],'Linewidth',2.5, 'Color',[0, 0.4470, 0.7410])
%plot(stateGrid, [mu_t2],'-.','LineWidth',2.5, 'Color',	[0.8500, 0.3250, 0.0980])
plot(stateGrid, [mu_t5],'-.','LineWidth',2.5, 'Color',	[1, 0, 1])
plot(stateGrid, ctilde,'k:','Linewidth',1.5)
scatter(ya_hist5, eta_hist,150,'kx')
%scatter(ya_hist5(2), eta_hist(2),150,'bx')
%scatter(ya_hist5(3), eta_hist(3),150,'cx')
%vline( mean(ya_hist), ':k', '', 16,1,2.5)
scatter([ya_hist(2)], [mu_t2(301)],150,'b','filled')
scatter([ya_hist(1)], [mu_t1(501)],150,'k','filled')
scatter(stateGrid(801),  mu_t5(801),150,'filled','MarkerFaceColor',[1 0 1])
scatter( stateGrid([361,451]), mu_t2([361,451]),150,'r','filled')
scatter([ya_hist(1)], [mu_t1(501)],150,'k','filled')
%vline( ya_hist(1), ':k', '', 16, 1, 1)
vline( ya_hist(1), ':k', '', 16, 1, 2.5)
vline( ya_hist(2), ':k', '', 16, 1, 2.5)
vline( ya_hist5(3), ':k', '', 16, 1, 2.5)
vline( stateGrid(361),   ':k', '', 16,1,2.5)
vline( stateGrid(451),   ':k', '', 16,1,2.5)
vline( stateGrid(801), ':k', '', 16,1,2.5)
hold off
xlim([1.25,4.5])
ylim([1.2, 1.435])
xticks( sort([ya_hist5;  stateGrid([361;451])]))
xticklabels({'$y_2$','$y_4$','$y_3$','$y_1$','$y_5$'})
yticks({})
%yticklabels( {'\sigma_c^2'})
ax = gca;
ax.FontSize = 17;
set(gca,'TickLabelInterpreter', 'latex')
ylabel('Action ($c$)','interpreter','latex')
xlabel('State ($y$)','interpreter','latex')
legend('$\widehat c_1(y)$', ' $\widehat c_2(y) = \widehat c_3(y)= \widehat c_4(y)$', '$\widehat c_5(y)$','$c^*(y)$','$\eta_t$','Location','Southeast','interpreter','latex','FontSize',17)
print -depsc example_chat3_v3.eps

%%  Figure of example path of choices c_t
figure
plot((1:5)', [mu_t1(501), mu_t2(301), mu_t2(451), mu_t2(361), mu_t5(701)]','k','Linewidth',1.5); hold on
plot((1:5)', ctilde([501,301,451,361,701]),'k:','Linewidth',1.5);
xlim([0.5,5.5])
ylim( [1.13,1.42])
%plot(0:5, zeros(1,6),'k--')
scatter( 1, mu_t1(501),125,'k','filled')
scatter( 2, mu_t2(301),125,'b','filled')
scatter( 3, mu_t2(451),125,'r','filled')
scatter( 4, mu_t2(361),125,'r','filled')
scatter( 5,  mu_t5(701),125,'filled','MarkerFaceColor',[1 0 1])
scatter( 1, ctilde(501),125,'k')
scatter( 2, ctilde(301),125,'b')
scatter( 3, ctilde(451),125,'r')
scatter( 4, ctilde(361),125,'r')
scatter( 5, ctilde(701),125,'MarkerEdgeColor',[1 0 1])
xticks( [1:5] )
yticks([])
yticklabels( {})
ax = gca;
ax.FontSize = 17;
set(gca,'TickLabelInterpreter', 'latex')
ylabel('Action ($c_t$)','interpreter','latex')
xlabel('Time period ($t$)','interpreter','latex')
hold off
legend('$c_t$','Location','Northwest','interpreter','latex','FontSize',17)
print -depsc example_ct_v3.eps
